transformer_engine/pytorch/csrc/extensions/cast.cpp
The alignment check split_section % 128 == 0 for fusion eligibility should be documented or extracted as a constant. This magic number appears in multiple places (quantizer.cpp line 1552 uses different values) and the inconsistency suggests potential bugs. Consider defining constexpr size_t NVFP4_FUSION_ROW_ALIGNMENT = 128; to centralize this requirement.

transformer_engine/pytorch/csrc/quantizer.cpp
When need_separate_columnwise_rng is true, a new RNG state is generated for columnwise quantization. However, the philox generator state is advanced for each call, which means the random sequences will differ between runs if the fusion eligibility changes (e.g., due to shape variations). For reproducibility across different code paths, consider using a deterministic offset based on tensor properties rather than sequential generator advancement.

transformer_engine/pytorch/csrc/quantizer.cpp, line 2292
Unnecessary intermediate buffer allocation when columnwise_usage=false

In the refactored unfused path, rht_output_t (a transposed buffer of shape cols × rows) is allocated unconditionally, even when columnwise_usage=false. In the old code this allocation was guarded inside if (columnwise_usage). The helper function itself skips the columnwise path when columnwise_usage=false, so this buffer is allocated but never used.

For large models (e.g. hidden_dim=4096, seq_len=2048, BF16) this wastes ~16MB of GPU memory per call and incurs unnecessary allocation overhead.

} else {
  auto& columnwise_quant_config_to_use =
      need_separate_columnwise_rng ? quant_config_columnwise : quant_config;
  // Only allocate the RHT transposed buffer if columnwise quantization is needed
  at::Tensor rht_output_t;
  TensorWrapper rht_output_t_cpp;
  if (columnwise_usage) {
    rht_output_t =
        allocateTorchTensor(static_cast<int>(cols), static_cast<int>(rows), input.dtype());
    rht_output_t_cpp.set_rowwise_data(rht_output_t.data_ptr(), input.dtype(),
                                      std::vector<size_t>{cols, rows});
  }
  this->quantize_with_rht_unfused_helper(input, out, rht_output_t_cpp, quant_config,
                                         columnwise_quant_config_to_use, stream);
}

transformer_engine/pytorch/csrc/quantizer.cpp, line 2289
Spurious memory allocation in unfused path when columnwise_usage=false

rht_output_t (a [cols × rows] BF16 tensor) is allocated unconditionally at lines 2282-2283 every time the unfused path is taken, even when columnwise_usage=false.

Inside quantize_with_rht_unfused_helper, this buffer is only accessed inside if (columnwise_usage) (line 2058), so for rowwise-only quantization the allocation is wasted. In the original code, this tensor was allocated inside if (columnwise_usage).

Consider guarding the allocation:

at::Tensor rht_output_t;
TensorWrapper rht_output_t_cpp;
if (columnwise_usage) {
    rht_output_t =
        allocateTorchTensor(static_cast<int>(cols), static_cast<int>(rows), input.dtype());
    rht_output_t_cpp.set_rowwise_data(rht_output_t.data_ptr(), input.dtype(),
                                      std::vector<size_t>{cols, rows});
}
this->quantize_with_rht_unfused_helper(input, out, rht_output_t_cpp, quant_config,
                                       columnwise_quant_config_to_use, stream);

transformer_engine/pytorch/csrc/quantizer.cpp, line 2271
Fused kernel used for rowwise-only case, potentially regressing performance

When rowwise_usage=true && columnwise_usage=false && eligible_for_rht_cast_fusion, the new code calls the full Cutlass GEMM-based nvte_quantize_with_hadamard_transform fused kernel (line 2269), even though RHT computation is not needed.

The fused kernel carries significant overhead — TMA pipeline infrastructure, cluster-launch-control scheduler, 512-thread CTAs — when only rowwise quantization is required. For inference workloads that only request rowwise output, dispatching a heavier GEMM kernel is slower than the simpler nvte_quantize_v2 approach.

Consider short-circuiting to nvte_quantize_v2 for rowwise-only cases:

if (eligible_for_rht_cast_fusion && columnwise_usage) {
    // fused path: rowwise + RHT columnwise in one kernel
    NVTE_CHECK(this->rht_matrix.defined() && this->rht_matrix.numel() > 0,
               "RHT matrix is not available.");
    auto rht_matrix_nvte = makeTransformerEngineTensor(this->rht_matrix);
    NVTE_SCOPED_GIL_RELEASE({
      nvte_quantize_with_hadamard_transform(input.data(), out.data(), rht_matrix_nvte.data(),
                                            quant_config, stream);
    });
} else if (eligible_for_rht_cast_fusion && !columnwise_usage) {
    // rowwise only – simpler kernel is sufficient
    NVTE_SCOPED_GIL_RELEASE({ nvte_quantize_v2(input.data(), out.data(), quant_config, stream); });
} else {
    // unfused fallback path
    ...
}

transformer_engine/pytorch/csrc/quantizer.cpp, line 2257
NVTE_USE_FAST_MATH env-var read on every quantize_impl call

transformer_engine::getenv<bool>("NVTE_USE_FAST_MATH") is called inside quantize_impl (line 2253), which executes on every forward/backward pass. For high-throughput training loops, this adds a getenv syscall on every kernel dispatch.

Since the PR notes that fast-math is "strongly recommended for performance", consider caching this value at construction time in the NVFP4Quantizer class or at module load time rather than re-reading it on every invocation.

transformer_engine/pytorch/csrc/quantizer.cpp, line 2288
Unnecessary buffer allocation when columnwise_usage=false

In the unfused path, rht_output_t (a cols × rows GPU tensor) is now allocated unconditionally before calling quantize_with_rht_unfused_helper, even when columnwise_usage=false (rowwise-only quantization). Inside the helper, rht_output_t_cpp is only used inside the if (columnwise_usage) block, so the allocation is wasted for the rowwise-only case.

Before this refactoring the allocation was correctly guarded inside if (columnwise_usage). The regression means every odd-shaped input with RHT in rowwise-only mode now incurs an unnecessary GPU memory allocation of size cols × rows × sizeof(BF16).

      // unfused path also needs memory allocation for intermediate buffer for RHT output
      at::Tensor rht_output_t;  // The RHT(x_t) output, in columnwise layout
      // This wrapper is going to be passed as input to the quantization kernel.
      TensorWrapper rht_output_t_cpp;  // Wrapper to contain the RHT(x) and RHT(x_t) outputs
      if (columnwise_usage) {
        rht_output_t =
            allocateTorchTensor(static_cast<int>(cols), static_cast<int>(rows), input.dtype());
        // NOTE (frsun): This is non-intuitive, we are writing the
        // result of transposed RHT to the output of rowwise.
        rht_output_t_cpp.set_rowwise_data(rht_output_t.data_ptr(), input.dtype(),
                                          std::vector<size_t>{cols, rows});
      }